#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <signal.h>
#include <stdbool.h>
#include <sys/time.h>
#include <json-c/json.h>
#include <unistd.h>  // For sleep function
#include "SerialControl_c.h"
#include <math.h>
#include <limits.h>  // For UINT_MAX
// Data structure to store frame and time information
typedef struct {
    unsigned int frame_index;
    double timestamp;  // Timestamp in seconds since epoch
} FrameRecord;

// Data structure for arm poses
typedef struct {
    double rads[6];  // Store 6 joint angles
} ArmPose;

// Global variables
static FrameRecord* frame_records = NULL;
static size_t frame_count = 0;
static size_t frame_capacity = 0;
static void* serial_handle = NULL;
static bool motion_started = false;
static double motion_start_time = 0.0;
static double motion_end_time = 0.0;
static ArmPose current_pose;
static ArmPose target_pose;
static bool capture_active = true;
static unsigned int target_frame_for_motion = 10;  // Start motion at 10th frame
static double motion_duration = 1.0;  // Default motion duration in seconds
static char episode_folder[256];

// Forward declarations
double get_current_time(void);
void add_frame_record(unsigned int index);
void check_and_start_motion(unsigned int frame_index);
void get_user_target_pose(void);
double calculate_motion_duration(ArmPose *current, ArmPose *target);

// Initialize controller
void init_capture_controller(void) {
    // Create episode folder with timestamp
    time_t now = time(NULL);
    struct tm *t = localtime(&now);
    snprintf(episode_folder, sizeof(episode_folder), 
            "/root/aixiaoxiao/armTrain/episode_%04d%02d%02d_%02d%02d%02d",
            t->tm_year + 1900, t->tm_mon + 1, t->tm_mday,
            t->tm_hour, t->tm_min, t->tm_sec);
    
    // Create episode directory
    mkdir(episode_folder, 0777);
    
    // Initialize storage
    frame_capacity = 1000;  // Initial capacity
    frame_records = (FrameRecord*)malloc(frame_capacity * sizeof(FrameRecord));
    if (!frame_records) {
        fprintf(stderr, "Failed to allocate memory for frame records\n");
        exit(EXIT_FAILURE);
    }
    
    // Initialize serial connection
    serial_handle = serial_control_init("/dev/ttyUSB0", 115200, 1000);
    if (!serial_handle) {
        fprintf(stderr, "Failed to initialize serial connection\n");
        free(frame_records);
        exit(EXIT_FAILURE);
    }
    
    serial_control_release_force(serial_handle, "arm");
    printf("INITIAL PLACE. ARM RELEASED \nPress 'l' to lock position: ");
    
    // Wait for user to press 'l' to lock
    char input;
    scanf(" %c", &input);
    while(getchar() != '\n'); // Clear input buffer
    
    if (input == 'l' || input == 'L') {
        get_new_target();
    }
}

// Calculate motion duration based on max angle difference
double calculate_motion_duration(ArmPose *current, ArmPose *target) {
    double max_diff = 0.0;
    
    for (int i = 0; i < 6; i++) {
        double diff = fabs(target->rads[i] - current->rads[i]);
        if (diff > max_diff) {
            max_diff = diff;
        }
    }
    
    // Linear scale: 0.5 radians -> 5.0 seconds, minimum 1 seconds
    double duration = max_diff * 10;
    if (duration < 1) duration = 1;
    if (duration > 9.9) duration = 9.9;
    return duration;
}

// Get target pose from user input
void get_user_target_pose(void) {
    printf("\nEnter target arm position (in radians):\n");
    
    // For each joint, get input from user
    for (int i = 0; i < 6; i++) {
        char input[64];
        double value;
        
        printf("Joint %d [current: %.4f rad]: ", i+1, current_pose.rads[i]);
        if (fgets(input, sizeof(input), stdin)) {
            // Remove newline
            input[strcspn(input, "\n")] = 0;
            
            // If empty, use current value
            if (strlen(input) == 0) {
                target_pose.rads[i] = current_pose.rads[i];
                printf("Using current value: %.4f rad\n", current_pose.rads[i]);
            } else {
                // Parse input
                if (sscanf(input, "%lf", &value) == 1) {
                    target_pose.rads[i] = value;
                    printf("Set to: %.4f rad\n", value);
                } else {
                    printf("Invalid input, using current value: %.4f rad\n", current_pose.rads[i]);
                    target_pose.rads[i] = current_pose.rads[i];
                }
            }
        } else {
            // Error in input, use current value
            target_pose.rads[i] = current_pose.rads[i];
            printf("Input error, using current value: %.4f rad\n", current_pose.rads[i]);
        }
    }

    printf("\nTarget arm position set to:\n");
    for (int i = 0; i < 6; i++) {
        printf("Joint %d: %.4f rad\n", i+1, target_pose.rads[i]);
    }
    
    printf("\nPress Enter to start the capture process...");
    getchar(); // Wait for user to press Enter
}

// Function to create a new session folder and copy PNG files
void create_new_session_cp() {
    // Create timestamped folder and move pngs
    time_t now = time(NULL);
    struct tm *t = localtime(&now);
    char session_folder[256];
    snprintf(session_folder, sizeof(session_folder), 
            "%s/session_%04d%02d%02d_%02d%02d%02d",
            episode_folder,
            t->tm_year + 1900, t->tm_mon + 1, t->tm_mday,
            t->tm_hour, t->tm_min, t->tm_sec);
    
    // Create session directory
    mkdir(session_folder, 0777);
    
    // Create pngs subdirectory in session folder
    char session_pngs[512];
    snprintf(session_pngs, sizeof(session_pngs), "%s/pngs", session_folder);
    mkdir(session_pngs, 0777);
    
    // Copy png files
    char cmd[1024];
    snprintf(cmd, sizeof(cmd), "cp /root/aixiaoxiao/armTrain/pngs/*.png %s", session_pngs);
    system(cmd);
        // Remove the first three PNG files
    // Remove the first three PNG files with the smallest %d in color_data_%d.png
    
    
    // Count the number of PNG files
    snprintf(cmd, sizeof(cmd), "ls /root/aixiaoxiao/armTrain/pngs/*.png | wc -l");
    FILE *png_count_pipe = popen(cmd, "r");
    int num_of_pngs = 0;
    if (png_count_pipe) {
        fscanf(png_count_pipe, "%d", &num_of_pngs);
        pclose(png_count_pipe);
    }
    
    // Linear interpolation of the arm poses
    json_object *motion_data = json_object_new_object();
    json_object *initial_pos = json_object_new_array();
    json_object *target_pos = json_object_new_array();
    json_object *data = json_object_new_array();
    
    for (int i = 0; i < 6; i++) {
        json_object_array_add(initial_pos, json_object_new_double(current_pose.rads[i]));
        json_object_array_add(target_pos, json_object_new_double(target_pose.rads[i]));
    }
    
    for (int i = 0; i < num_of_pngs; i++) {
        json_object *pose = json_object_new_array();
        for (int j = 0; j < 6; j++) {
            double interpolated_value = current_pose.rads[j] + (target_pose.rads[j] - current_pose.rads[j]) * (i / (double)(num_of_pngs - 1));
            json_object_array_add(pose, json_object_new_double(interpolated_value));
        }
        json_object_array_add(data, pose);
    }
    
    json_object_object_add(motion_data, "numOfMotion", json_object_new_int(num_of_pngs));
    json_object_object_add(motion_data, "initial_pos", initial_pos);
    json_object_object_add(motion_data, "target_pos", target_pos);
    json_object_object_add(motion_data, "motion_start_time", json_object_new_double(motion_start_time));
    json_object_object_add(motion_data, "motion_end_time", json_object_new_double(motion_end_time));
    json_object_object_add(motion_data, "data", data);
    
    // Write JSON to file
    char json_file_path[512];
    snprintf(json_file_path, sizeof(json_file_path), "%s/motion_data.json", session_folder);
    FILE *json_file = fopen(json_file_path, "w");
    if (json_file) {
        fprintf(json_file, "%s", json_object_to_json_string_ext(motion_data, JSON_C_TO_STRING_PRETTY));
        fclose(json_file);
    }
    
    // Clear original pngs directory
    snprintf(cmd, sizeof(cmd), "rm -f /root/aixiaoxiao/armTrain/pngs/*.png");
    system(cmd);
    
    // Clean up JSON object
    json_object_put(motion_data);
}

// Function to get new target pose and recalculate motion duration
void get_new_target() {
    // Update current pose before getting new target
    if (serial_control_get_rads_data(serial_handle, "arm", current_pose.rads, 6) <= 0) {
        fprintf(stderr, "Failed to read current arm pose\n");
        // If read fails, keep existing values
    } else {
        printf("\nUpdated current arm position (radians):\n");
        for (int i = 0; i < 6; i++) {
            printf("Joint %d: %.4f rad\n", i+1, current_pose.rads[i]);
        }
    }
    
    // Release force to allow manual movement
    serial_control_release_force(serial_handle, "arm");
    printf("NOW YOU CAN MOVE THE ARM TO THE NEXT POS AND HOLD\nPress 'l' to lock position: ");
    
    // Wait for user to press 'l' to lock
    char input;
    scanf(" %c", &input);
    while(getchar() != '\n'); // Clear input buffer
    
    if (input == 'l' || input == 'L') {
        // Recover force to lock arm in position
        serial_control_recover_force(serial_handle, "arm");
        printf("Position locked!\n");
        
        // Get the current position as target position
        if (serial_control_get_rads_data(serial_handle, "arm", target_pose.rads, 6) <= 0) {
            fprintf(stderr, "Failed to read target arm pose\n");
            // If read fails, keep existing values
        } else {
            printf("\nNew target arm position (radians):\n");
            for (int i = 0; i < 6; i++) {
                printf("Joint %d: %.4f rad\n", i+1, target_pose.rads[i]);
            }
        }
        
        // Drive arm back to the starting position as fast as possible (1 second)
        printf("Returning to starting position...\n");
        serial_control_set_move_time(serial_handle, 1000); // 1 second
        serial_control_send_joint_command(serial_handle, current_pose.rads, 6, "arm");
        
        // Sleep for 1 second while arm moves
        sleep(1);
    }
    
    // Recalculate motion duration based on the new target
    motion_duration = calculate_motion_duration(&current_pose, &target_pose);
    printf("New motion duration calculated: %.2f seconds\n", motion_duration);
    sleep(1);
}

// Record a new frame
void record_frame(unsigned int index) {
    // Store frame record with current timestamp
    add_frame_record(index);
    
    // Check if we should start motion
    check_and_start_motion(index);
    
    // Check if motion has ended
    if (motion_started && get_current_time() > motion_end_time) {
        double elapsed_ms = (get_current_time() - motion_start_time) * 1000.0;
        printf("\nMotion ended in %.2f ms\n", elapsed_ms);
        
        // Reset motion flags
        motion_started = false;
        
        // Create new session and copy PNG files
        create_new_session_cp();
        
        // Get new target pose and recalculate motion duration
        get_new_target();
    }
}

// Add a frame record to our storage
void add_frame_record(unsigned int index) {
    // Resize if needed
    if (frame_count >= frame_capacity) {
        frame_capacity *= 2;
        FrameRecord* new_records = (FrameRecord*)realloc(frame_records, 
                                                       frame_capacity * sizeof(FrameRecord));
        if (!new_records) {
            fprintf(stderr, "Failed to resize frame records array\n");
            return;
        }
        frame_records = new_records;
    }
    
    // Add new record
    frame_records[frame_count].frame_index = index;
    frame_records[frame_count].timestamp = get_current_time();
    frame_count++;
}

// Get current time in seconds with high precision
double get_current_time(void) {
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return (double)tv.tv_sec + (double)tv.tv_usec / 1000000.0;
}

// Check if we should start arm motion
void check_and_start_motion(unsigned int frame_index) {
    if (!motion_started) {
        printf("Starting arm motion at frame %u\n", frame_index);
        
        // Record motion start time
        motion_start_time = get_current_time();
        
        // Set motion duration in milliseconds for the serial command
        int move_time_ms = (int)(motion_duration * 1000.0);
        printf("Setting move time to %d ms\n", move_time_ms);
        serial_control_set_move_time(serial_handle, move_time_ms);
        
        // Send command to move arm
        serial_control_send_joint_command(serial_handle, target_pose.rads, 6, "arm");
        
        // Set motion end time based on calculated duration
        motion_end_time = motion_start_time + motion_duration;
        
        // Mark motion as started
        motion_started = true;
        
    }
}

